# Dockerfile
# 使用一个包含 PyTorch 和 CUDA 的基础镜像
# 你可以从 NVIDIA NGC (nvcr.io/nvidia/pytorch) 或 Docker Hub (pytorch/pytorch) 选择
# 确保 PyTorch 版本与 CUDA 版本兼容你的 GPU 和驱动
# FROM pytorch/pytorch:2.1.0-cuda11.8-cudnn8-runtime
# FROM nvcr.io/nvidia/pytorch:24.01-py3
FROM registry.deepbay.ai/ap-dp-deepbay/pytorch:2.2.0

# 设置工作目录
WORKDIR /app

# 安装依赖
# transformers, datasets, accelerate 用于 Hugging Face 模型和训练
# sentencepiece 通常是某些 tokenizer 的依赖
# Git 用于某些模型可能需要从源拉取代码 (如果 trust_remote_code=True)
RUN pip install --no-cache-dir \
    transformers==4.38.2 \
    datasets==2.18.0 \
    accelerate==0.27.2 \
    sentencepiece==0.2.0 \
    tensorboard \
    # 如果你的模型需要特定版本的 protobuf
    # protobuf==3.20.0

# 复制脚本到工作目录
COPY train.py .
COPY prepare_data.py .

# 设置容器启动时执行的命令
# ENTRYPOINT ["python", "train.py"]
# 使用 CMD 允许在 `docker run` 或 Kubernetes YAML 中覆盖参数
# 例如，你可以通过 Kubernetes Job 的 args 来传递不同的训练参数
# 这里我们不设置 ENTRYPOINT 或 CMD，因为 PyTorchJob 会指定 command/args
# 如果你希望镜像本身是可执行的（例如 `docker run myimage --num_train_epochs 2`），
# 你可以取消注释 ENTRYPOINT 或添加 CMD ["python", "train.py", "--help"]

# 确保输出目录存在，如果脚本不创建它的话
RUN mkdir -p /app/model_output

# 声明 Python 脚本为入口点，这样 PyTorchJob 可以直接调用它
# 如果你的脚本需要通过 `torchrun` 启动以进行分布式训练，你需要相应地调整
# 对于单节点训练，直接运行 python 脚本即可
# ENTRYPOINT ["python", "/app/train.py"]